package com.farm2.llm.client.impl;

import com.alibaba.fastjson2.JSON;
import com.farm2.llm.client.inter.Farm2EmbedClientInter;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.farm2.llm.domain.LlmClient;
import org.farm2.tools.json.FarmJsons;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Farm2EmbedAliyunClientImpl implements Farm2EmbedClientInter {
    private LlmClient paras;

    @Override
    public String getName() {
        return "Aliyun_Embedding_SDK";
    }

    @Override
    public void init(LlmClient client) {
        paras = client;
    }

    @Override
    public float[] getEmbedding(String text) {
        HttpClient client = HttpClient.newHttpClient();
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", paras.getModelkey());
        requestBody.put("input", text);
        requestBody.put("dimensions", paras.getDimensions());
        requestBody.put("encoding_format", paras.getEnformat());
        try {
            String requestBodyString = FarmJsons.toJson(requestBody);
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(paras.getBaseurl()))
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + paras.getApikey())
                    .POST(HttpRequest.BodyPublishers.ofString(requestBodyString))
                    .build();
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            if (response.statusCode() == 200) {
                float[] embeddings = getEmbeddings(response.body());
                return embeddings;
            } else {
                throw new RuntimeException("Failed to retrieve response, status code: " + response.statusCode() + ", response: " + response.body());
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static void main(String[] args) {
        String apiKey = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxx";
        String baseUrl = "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings";
        HttpClient client = HttpClient.newHttpClient();

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "text-embedding-v3");
        requestBody.put("input", "风急天高猿啸哀，渚清沙白鸟飞回，无边落木萧萧下，不尽长江滚滚来");
        requestBody.put("dimensions", 1024);
        requestBody.put("encoding_format", "float");

        try {
            String requestBodyString = FarmJsons.toJson(requestBody);
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(baseUrl))
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + apiKey)
                    .POST(HttpRequest.BodyPublishers.ofString(requestBodyString))
                    .build();

            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            if (response.statusCode() == 200) {
                float[] embeddings = getEmbeddings(response.body());
                System.out.println("Response: " + embeddings);
            } else {
                System.out.printf("Failed to retrieve response, status code: %d, response: %s%n", response.statusCode(), response.body());
            }
        } catch (Exception e) {
            System.err.println("Error: " + e.getMessage());
        }
    }

    public static float[] getEmbeddings(String json) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {

            EmbeddingResponse response = objectMapper.readValue(json, EmbeddingResponse.class);

            List<Double> embeddingList = response.getData().get(0).getEmbedding();

            float[] embeddingArray = new float[embeddingList.size()];
            for (int i = 0; i < embeddingList.size(); i++) {
                embeddingArray[i] = embeddingList.get(i).floatValue();
            }
            return embeddingArray;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static class EmbeddingResponse {

        @JsonProperty("data")
        private List<EmbeddingData> data;

        @JsonProperty("object") // 添加对 "object" 字段的支持
        private String object;

        @JsonProperty("model") // 添加对 "model" 字段的支持
        private String model;

        @JsonProperty("usage") // 添加对 "usage" 字段的支持
        private Usage usage;

        @JsonProperty("id") // 添加对 "id" 字段的支持
        private String id;

        public List<EmbeddingData> getData() {
            return data;
        }

        public void setData(List<EmbeddingData> data) {
            this.data = data;
        }

        public String getObject() { // 添加 getter 方法
            return object;
        }

        public void setObject(String object) { // 添加 setter 方法
            this.object = object;
        }

        public String getModel() { // 添加 getter 方法
            return model;
        }

        public void setModel(String model) { // 添加 setter 方法
            this.model = model;
        }

        public Usage getUsage() { // 添加 getter 方法
            return usage;
        }

        public void setUsage(Usage usage) { // 添加 setter 方法
            this.usage = usage;
        }

        public String getId() { // 添加 getter 方法
            return id;
        }

        public void setId(String id) { // 添加 setter 方法
            this.id = id;
        }

        public static class EmbeddingData {
            @JsonProperty("embedding")
            private List<Double> embedding;

            @JsonProperty("index")
            private int index;

            @JsonProperty("object")
            private String object;

            public List<Double> getEmbedding() {
                return embedding;
            }

            public void setEmbedding(List<Double> embedding) {
                this.embedding = embedding;
            }

            public int getIndex() {
                return index;
            }

            public void setIndex(int index) {
                this.index = index;
            }

            public String getObject() {
                return object;
            }

            public void setObject(String object) {
                this.object = object;
            }
        }

        public static class Usage {
            @JsonProperty("prompt_tokens")
            private int promptTokens;

            @JsonProperty("total_tokens")
            private int totalTokens;

            public int getPromptTokens() {
                return promptTokens;
            }

            public void setPromptTokens(int promptTokens) {
                this.promptTokens = promptTokens;
            }

            public int getTotalTokens() {
                return totalTokens;
            }

            public void setTotalTokens(int totalTokens) {
                this.totalTokens = totalTokens;
            }
        }
    }
}
